#! /usr/bin/env python
# -*- coding: UTF-8 -*-
'''修改一个python项目的包引用方式

    1. 绝对路径包引用改为相对路径包引用
'''
import os
import sys
from pathlib import Path
import pdb
import logging

# 设置包搜索路径
__p = Path(__file__).parent
sys.path.append(str(__p))

from py_lib.logger_utils import logger_init
from py_lib.my_exception import MyException
from py_lib.python_import_fix.args_utils import args_parse

def __is_top_file(p):
    '''p是不是入口函数
    
        对不同的项目，这个函数需要定制。目前这个适用于onnx-quantization
    '''
    topFiles = ['frontends/caffe/caffe2onnx.py'
            , 'frontends/paddle/paddle2onnx.py'
            , 'frontends/pytorch/pytorch2onnx.py'
            , 'frontends/tensorflow/tensorflow2onnx.py'
            , 'test/st/run_test.py'
            , 'test/frontends/caffe/test_caffe2onnx.py'
            , 'test/smoke/smoke_onnx.py'
            , 'test/smoke/smoke_test.py']
    selfPkg = __file_pkg_name(p)
    if selfPkg is None: 
        # run_quantization.py等 
        return True
    for s in topFiles:
        if p.match(s):
            return True
    # tools/xxxx.py都是顶层文件
    if str(p.parent) == 'tools':
        return True
    return False

def __file_pkg_name(p):
    '''返回python文件所在的完整包名
    
       返回: 在p是顶层文件时返回None
    '''
    assert p.exists() and p.is_file(), '{p} not found or is not a file'
    dParts = p.parent.parts
    # 忽略顶层python文件，因为它们不在任何包内
    if len(dParts) == 0:
        return None
    return '.'.join(dParts)

def __build_pkg_names(files):
    '''得到所有的包名, 包名按从深到浅排序.
    
        包名从深到浅排序是为了不发生重复替换 
    '''
    pkgs = []
    for p in files:
        if not p.exists() or not p.is_file():
            continue
        pkg = __file_pkg_name(p)
        if pkg is None:
            # p是顶层文件
            continue
        if pkg not in pkgs:
            pkgs.append(pkg)
    # 包名排序
    pkgs.sort(key = lambda x: len(x.split('.')), reverse = True)
    return pkgs

def __build_dir_py_file_list(d):
    '''建造python文件列表, 列表的每个元素是一个Path对象
    
        d: str or Path对象, 要搜索的目录
    '''

    files = [] 
    for f in list(os.popen(f'find {d} -name "*.py" -type f')):
        files.append(Path(f.strip()))
    return files

def __build_py_file_list(paths):
    '''建造python文件列表, 列表的每个元素是一个Path对象
    
        paths: 是字符列表
    '''
    files = []
    for s in paths:
        p = Path(s)
        if not p.exists():
            raise MyException(f'{p} not found')
        if p.is_file():
            files.append(p)
        elif p.is_dir():
            files.extend(__build_dir_py_file_list(p)) 

    return files

def __do_sed_cmd(cmd):
    logger.debug(cmd)
    r = os.system(cmd)
    assert r == 0, f'执行命令失败:{cmd} 错误码:{r}'

def __replace_from_import(p, impPkg, repPkg, impRest):
    # 组织sed命令
    # refPkg全是...
    if not impRest:
        cmd = f'sed -i \'s/^\\([ \\t]*from\\)[ \\t]*{impPkg}\\./\\1 {repPkg}/\' \"{p}\"'
        __do_sed_cmd(cmd)
        cmd = f'sed -i \'s/^\\([ \\t]*from\\)[ \\t]*{impPkg}\\([ \\t]\\)/\\1 {repPkg}\\2/\' \"{p}\"'
        __do_sed_cmd(cmd)
    else:
        cmd = f'sed -i \'s/^\\([ \\t]*from\\)[ \\t]*{impPkg}\\([. \\t]\\)/\\1 {repPkg}\\2/\' \"{p}\"'
        __do_sed_cmd(cmd)

def __replace_real_import(p, impPkg):
    cmd = f'sed -i \'s/^\\([ \\t]*import\\)[ \\t]\\+\\({impPkg}\\)/\\1 {args.top_pkg}.\\2/\' \"{p}\"'
    __do_sed_cmd(cmd)

def __replace_import(p, selfPkg, impPkg):
    '''将文件p中的对impPkg的import全部改为相对路径

        p: 待修改文件
        selfPkg：待修改文件所在的包
        impPkg: 待修改的包
    '''
    # if str(p) == 'modules/mask.py' and impPkg == 'modules':
    #    pdb.set_trace()
    selfParts = selfPkg.split('.')
    impParts = impPkg.split('.')
    # p所在包的深度
    selfDepth = len(selfParts)
    # 比较并计算selfPkg和impPkg不同部分的index
    # 待替换包是否还有剩下的
    impRest = False
    for i in range(len(impParts)):
        if i >= selfDepth:
            if i < len(impParts):
                impRest = True
            break
        if selfParts[i] != impParts[i]:
            impRest = True
            break
    idx = i
    logger.debug(f'self depth: {selfDepth}, idx: {idx}, impRest: {impRest}')
    # 需要的点数,至少需要一个点
    dotN = (selfDepth - idx) 
    # 替换包名, idx是不可能超出impParts的range的
    if not impRest:
        repPkg = '.' * dotN 
    else:
        repPkg = '.' * dotN  + '.' + '.'.join(impParts[idx:])
    logger.debug(f'{p}, self pkg: {selfPkg}, import pkg: {impPkg} --> replace pkg: {repPkg}')
    # from xxxx import yyyy形式
    __replace_from_import(p, impPkg, repPkg, impRest)
    # import xxxx [as] xxxx形式 
    __replace_real_import(p, impPkg)

def __fix_at_end(p):
    '''修复可能出现的`from[空格]import`问题'''
    cmd = f'sed -i \'s/^\\([ \\t]*from\\)[ \\t]*import /\\1 . import /\' \"{p}\"'
    __do_sed_cmd(cmd)

def __replace_top_file(p, impPkg):
    '''替换顶层文件的Import
    
        impPkg: 待处理的import包名
    '''
    if args.top_pkg is None:
        # 没有指定顶层包名，不做顶层文件替换
        return

    pkgParts = impPkg.split('.')
    if len(pkgParts) == 0:
        return
    # 如果要替换的包的第一层已经是指定的包名，则不用替换
    if args.top_pkg == pkgParts[0]:
        return

    logger.debug(f'Replace import for {p}')
    for prefix in ['from', 'import']:
        cmd = f'sed -i \'s/^\\([ \\t]*{prefix}\\)[ \\t]\\+\\({impPkg}\\)/\\1 {args.top_pkg}.\\2/\' \"{p}\"'
        __do_sed_cmd(cmd)

def __process_file(p, pkgs):
    '''将文件p的import改为相对路径'''
    print(f'{p}...')
    if not p.exists or not p.is_file():
        raise MyException(f'{p} not found or is not a file')
    selfPkg = __file_pkg_name(p)
    if args.pkg is None:
        # 一个包一个包的处理
        for impPkg in pkgs:
            # pdb.set_trace()
            if __is_top_file(p):
                # 处理顶层文件
                __replace_top_file(p, impPkg)
                continue
            __replace_import(p, selfPkg, impPkg)
    else:
        # 仅处理指定包
        if __is_top_file(p):
            # 处理顶层文件
            __replace_top_file(p, impPkg)
        else:
            __replace_import(p, selfPkg, args.pkg)

    # 修复可能出现的`from[空格]import`问题
    __fix_at_end(p)
    print(f'Done')

def main():
    # 分析命令行参数
    global args
    args = args_parse()
    # 初始化logger
    logger_init(args.verbosity)
    global logger
    logger = logging.getLogger(__name__)
    # 建造python文件列表, 列表的每个元素是一个Path对象
    files = __build_dir_py_file_list('.')
    logger.debug(f'Total {len(files)} pythons files found')
    # 得到所有的包名
    pkgs = __build_pkg_names(files)
    logger.debug(f'pkg names: {pkgs}')

    if len(args.path) == 0:
        # 逐个处理python文件
        for f in files:
            __process_file(f, pkgs)
    else:
        files = __build_py_file_list(args.path)
        for f in files:
            __process_file(Path(f), pkgs)

if __name__ == '__main__':
    try:
        main()
    except MyException as e:
        print(e)
